import json
import chardet
import matplotlib as mpl
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
from datetime import datetime
from matplotlib import ticker
from matplotlib.lines import Line2D
from adjustText import adjust_text

# 全局配置解决字体和性能问题
mpl.rcParams.update({
    'font.family': 'Microsoft YaHei',
    'axes.unicode_minus': False,
    'agg.path.chunksize': 20000
})

# 文件路径配置
kline_trade_file = 'D:\\work\\code\\clark\\gitee\\big_a\\datas\\backtest\\sz002468\\20250529-20250715_chart.txt'

def detect_encoding(file_path):
    """检测文件编码"""
    with open(file_path, 'rb') as f:
        return chardet.detect(f.read(10000))['encoding']

def load_kline_data(file_path):
    """加载K线数据"""
    kline_data = []
    encoding = detect_encoding(file_path)
    with open(file_path, 'r', encoding=encoding, errors='replace') as file:
        for line in file:
            if line.strip().startswith('{"add":'):
                try:
                    kline_data.append(json.loads(line))
                except:
                    continue
    return kline_data

def load_trade_data(file_path):
    """加载交易数据"""
    trade_data = []
    encoding = detect_encoding(file_path)
    with open(file_path, 'r', encoding=encoding, errors='replace') as file:
        for line in file:
            if line.strip().startswith('{"cash":'):
                try:
                    trade_data.append(json.loads(line))
                except:
                    continue
    return trade_data

# 加载数据
kline_data = load_kline_data(kline_trade_file)
trade_data = load_trade_data(kline_trade_file)

# 准备K线数据
dates = [datetime.strptime(k['date'], '%Y-%m-%d %H:%M:%S') for k in kline_data]
opens = [k['open'] for k in kline_data]
highs = [k['high'] for k in kline_data]
lows = [k['low'] for k in kline_data]
closes = [k['close'] for k in kline_data]
volumes = [k['volume'] for k in kline_data]

# 创建图表（修复GridSpec行数匹配问题）
fig = plt.figure(figsize=(18, 12), layout="constrained")
gs = fig.add_gridspec(nrows=3, ncols=1, height_ratios=[4, 1, 1], hspace=0.05)  # 3行布局

# 1. K线图区域
ax1 = fig.add_subplot(gs[0])
ax1.set_title(f'资产价格走势与交易点 (20250529-20250715)', fontsize=14, fontweight='bold')
ax1.set_ylabel('价格', fontsize=12)
ax1.grid(True, linestyle='--', alpha=0.7)

# 动态计算K线宽度
time_intervals = np.diff(mdates.date2num(dates))
width = np.mean(time_intervals) * 0.6 if len(time_intervals) > 0 else 0.004

# 绘制K线（解决重叠问题）
for i in range(len(dates)):
    color = 'red' if closes[i] >= opens[i] else 'green'
    # 上下影线
    ax1.vlines(dates[i], lows[i], highs[i], color=color, linewidth=0.8)
    # 实体部分
    ax1.bar(dates[i],
            height=abs(closes[i]-opens[i]),
            bottom=min(opens[i], closes[i]),
            width=width,
            color=color,
            edgecolor=color)

# 2. 成交量区域
ax2 = fig.add_subplot(gs[1], sharex=ax1)
ax2.bar(dates, volumes, width=width*0.8,
        color=['red' if c >= o else 'green' for o, c in zip(opens, closes)])
ax2.set_ylabel('成交量', fontsize=12)
ax2.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{x/1e6:.1f}M'))
ax2.grid(True, linestyle='--', alpha=0.7)

# 3. 资产曲线区域
ax3 = fig.add_subplot(gs[2], sharex=ax1)
trade_dates = [datetime.strptime(t['date'], '%Y-%m-%d %H:%M:%S') for t in trade_data]
assets = [t['totalAsset'] for t in trade_data]
ax3.plot(trade_dates, assets, 'b-', linewidth=2)
ax3.set_ylabel('总资产', fontsize=12)
ax3.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'¥{x/1000:.0f}K'))
ax3.grid(True, linestyle='--', alpha=0.7)

# 设置时间刻度（精确到分钟）
days = (dates[-1] - dates[0]).days
interval = max(1, days // 10)  # 自动计算刻度间隔
ax1.xaxis.set_major_locator(mdates.DayLocator(interval=interval))
ax1.xaxis.set_minor_locator(mdates.HourLocator(byhour=[9, 11, 13, 15]))
ax1.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d %H:%M'))
plt.setp(ax1.get_xticklabels(), rotation=45, ha='right')

# 识别买卖点
buy_points, sell_points = [], []
for trade in trade_data:
    dt = datetime.strptime(trade['date'], '%Y-%m-%d %H:%M:%S')
    price = trade['price']
    operation = trade['operation']
    volume = trade.get('volume', trade.get('position', 0))

    # 查找最近的时间点
    min_idx = min(range(len(dates)), key=lambda i: abs(dates[i] - dt))
    if "买入" in operation:
        buy_points.append((dates[min_idx], price, volume, min_idx))
    elif "卖出" in operation:
        sell_points.append((dates[min_idx], price, volume, min_idx))

# 绘制买卖点标记（解决标注覆盖问题）
texts = []
for point in buy_points:
    date, price, volume, idx = point
    t = ax1.text(date, price, f'B {price:.2f}\n+{volume}',
                 fontsize=9, color='red', ha='center', va='top',
                 bbox=dict(boxstyle="round,pad=0.2", fc='white', ec='red', alpha=0.8))
    texts.append(t)

for point in sell_points:
    date, price, volume, idx = point
    t = ax1.text(date, price, f'S {price:.2f}\n-{volume}',
                 fontsize=9, color='green', ha='center', va='bottom',
                 bbox=dict(boxstyle="round,pad=0.2", fc='white', ec='green', alpha=0.8))
    texts.append(t)

# 智能调整标注位置（避免重叠）
adjust_text(texts,
            arrowprops=dict(arrowstyle='->', color='gray', lw=0.8, alpha=0.7, shrinkA=10),
            expand_points=(1.5, 1.5),
            expand_text=(1.2, 1.2),
            force_points=0.5,
            ax=ax1)

# 添加图例
legend_elements = [
    Line2D([0], [0], color='red', lw=4, label='阳线'),
    Line2D([0], [0], color='green', lw=4, label='阴线'),
    Line2D([0], [0], marker='o', color='red', label='买入点', markersize=8, linestyle='None'),
    Line2D([0], [0], marker='o', color='green', label='卖出点', markersize=8, linestyle='None')
]
ax1.legend(handles=legend_elements, loc='upper left', fontsize=9)

# 调整布局防止重叠
plt.subplots_adjust(
    left=0.06, right=0.95,
    top=0.95, bottom=0.1,
    hspace=0.15
)

plt.savefig('专业K线分析图.png', dpi=150, bbox_inches='tight')
plt.show()